import jieba
import joblib
import numpy as np
from sklearn.feature_extraction.text import TfidfVectorizer
from keras.models import load_model
from keras.preprocessing.sequence import pad_sequences

# 定义分词函数
def chinese_tokenizer(text):
    return list(jieba.cut(text))

# 加载训练好的模型和工具
print("正在加载模型和工具...")
lr_model = joblib.load('lr_model.pkl')  # Logistic Regression 模型
lstm_model = load_model('lstm_model_best.keras')  # LSTM 模型
vectorizer = joblib.load('tfidf_vectorizer.pkl')  # TF-IDF 向量化器
label_encoder = joblib.load('label_encoder.pkl')  # 标签编码器
print("模型和工具加载完成！")

# 定义预测函数
def predict_text(text, model_type='lr'):
    """
    根据输入的文本进行预测，并返回标签
    """
    # 文本向量化
    x_data_vectorized = vectorizer.transform([text]).toarray()

    if model_type == 'lr':
        # Logistic Regression 模型预测
        prediction = lr_model.predict(x_data_vectorized)
    elif model_type == 'lstm':
        # LSTM 模型预测
        max_length = 100  # 假设文本最大长度为100
        x_data_padded = pad_sequences(x_data_vectorized, padding='post', maxlen=max_length)
        prediction = lstm_model.predict(x_data_padded)
        prediction = np.argmax(prediction, axis=1)  # 获取概率最大的类别
    
    # 解码预测的标签
    predicted_label = label_encoder.inverse_transform(prediction)[0]
    return predicted_label

# 进入对话模式
print("\n欢迎进入文本分类预测对话模式！")
print("请输入需要分类的文本，输入“退出”即可结束程序。\n")

while True:
    user_input = input("请输入文本: ").strip()
    if user_input.lower() in ["退出", "exit"]:
        print("程序已退出，再见！")
        break
    try:
        # 选择模型类型（可根据需要更改为 'lstm'）
        model_type = 'lr'  
        predicted_label = predict_text(user_input, model_type=model_type)
        print(f"预测结果: {predicted_label}\n")
    except Exception as e:
        print(f"发生错误: {e}\n")
